from langchain_core.prompts import FewShotPromptTemplate, PromptTemplate
from langchain_openai import ChatOpenAI
import os
from dotenv import load_dotenv
from pydantic import SecretStr


# 加载.env文件中的环境变量
load_dotenv()

llm = ChatOpenAI(
    model="qwen-flash",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    api_key=SecretStr(os.getenv("DASHSCOPE_API_KEY")),
    streaming=True
)

# 示例模板
example_template = "输入:{input}\n输出:{output}"

# 示例
examples = [
    {"input": "将'你好'翻译成英文", "output": "Hello"},
    {"input": "将'好的'翻译成英文", "output": "Ok"}
]

# 少样本提示词模板
few_shot_prompt_template = FewShotPromptTemplate(
    examples=examples,
    example_prompt=PromptTemplate.from_template(example_template),
    prefix="请将下面的中文翻译为英文:",
    suffix="输入:{text}\n输出:",
    input_variables=["text"]
)

prompt = few_shot_prompt_template.format(text="我热爱编程")
print(prompt)

resp = llm.stream(prompt)
for chunk in resp:
    print(chunk.content, end="")
